package backend;

import backend.inst.MipsInst;
import backend.value.Offset;
import backend.value.Register;

import java.util.*;

/**
 * 临时寄存器池，用于管理目标机器的寄存器。
 *
 * @author 杨凯伟
 * @date 2024/10/30
 */
public class TemRegisterPool {
    private final Target targetCode;
    private final Queue<Register> usageQueue = new ArrayDeque<>();
    private final Stack<Register> unusedRegisters = new Stack<>();
    private final List<Register> availableRegisters = new ArrayList<>();
    private final Map<Offset, Register> addressToRegisterMap = new HashMap<>();
    private final Map<Register, Offset> registerToAddressMap = new HashMap<>();


    public TemRegisterPool(Target target, List<Register> registers) {
        this.targetCode = target;
        this.availableRegisters.addAll(registers);
        this.unusedRegisters.addAll(registers);
    }

    public Register allocTempRegister(Offset addr, boolean firstTime) {
        if (hasBeenAllocated(addr)) {// 已经分配了，更新使用状态
            return updateRegisterUsage(addr);
        }
        if (unusedRegisters.isEmpty()) {//全部都分配，移除最早分配的
            evictOldestRegister();
        }
        Register regToAlloc = allocateNewRegister();
        initializeRegister(regToAlloc, addr, firstTime);
        return regToAlloc;
    }

    private Register updateRegisterUsage(Offset addr) {
        Register reg = getRegister(addr);
        usageQueue.remove(reg);
        usageQueue.add(reg);
        return reg;
    }

    private void evictOldestRegister() {
        Register regToKill = usageQueue.poll();
        Offset regToKillAddr = registerToAddressMap.get(regToKill);
        addressToRegisterMap.remove(regToKillAddr);
        registerToAddressMap.remove(regToKill);
        targetCode.addCode(new MipsInst("sw", regToKill, regToKillAddr));
        unusedRegisters.add(regToKill);
    }

    private Register allocateNewRegister() {
        Register regToAlloc = unusedRegisters.pop();
        usageQueue.add(regToAlloc);
        return regToAlloc;
    }

    private void initializeRegister(Register regToAlloc, Offset addr, boolean firstTime) {
        if (!firstTime) {
            targetCode.addCode(new MipsInst("lw", regToAlloc, addr));
        }
        addressToRegisterMap.put(addr, regToAlloc);
        registerToAddressMap.put(regToAlloc, addr);
    }

    public void writeBackToMemoryForAll() {
        while (!usageQueue.isEmpty()) {
            var reg = usageQueue.poll();
            var addr = registerToAddressMap.get(reg);
            targetCode.addCode(new MipsInst("sw", reg, addr));
        }
    }

    public void reset() {
        unusedRegisters.clear();
        usageQueue.clear();
        addressToRegisterMap.clear();
        registerToAddressMap.clear();
        unusedRegisters.addAll(availableRegisters);
    }

    public boolean hasBeenAllocated(Offset addr) {
        return addressToRegisterMap.containsKey(addr);
    }

    public Register getRegister(Offset addr) {
        return addressToRegisterMap.get(addr);
    }
}
